import torch
import numpy as np
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader
from tqdm import trange, tqdm

from Datasets import LyricsDataset
seq_len = 48


def train():
    for epoch in trange(epochs, desc="Epoch",leave=False):
        total_loss = 0
        for x, y in tqdm(data_loader,desc="batch",leave=False):
            x = x.to(device)
            y = y.to(device)
            optimizer.zero_grad()
            output = model(x)
            # 修改以下行：使用 y 而不是 dataset.word2index 进行切片操作
            loss = criterion(output[:, :-1].reshape(-1, len(dataset.word2index)), y[:, 1:].reshape(-1))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(data_loader)}")

if __name__ == '__main__':
    dataset = LyricsDataset(seq_len=seq_len)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = torch.load("lyrics.pt", weights_only=False).to(device)
    # 添加batch_size参数
    data_loader = DataLoader(dataset, batch_size=256, shuffle=True)
    epochs = 10  # 增加训练轮数
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)  # 降低学习率
    criterion = nn.CrossEntropyLoss(ignore_index=dataset.word2index["<PAD>"])
    train()
    torch.save(model, "lyrics.pt")